import paddle
import pandas as pd
from net import QKS

voc_id = pd.read_pickle("voc_data.pandas_pickle")["data"]
voc = pd.read_pickle("src_data_token_and_padding.pandas_pickle")["data"]
voc = voc.head(1).values.tolist()[0]
print("".join([voc_id[i] for i in voc]))

net = QKS(len(voc_id), 256, 8)
net.load_dict(paddle.load("long_attention_model (3)"))
net.eval()
pre_in = paddle.to_tensor(voc[:-13]).reshape([1, -1]).astype("int64")
label_l = paddle.to_tensor(voc[1:]).reshape([1, -1]).astype("int64")
for i in range(13):
    outl,sl=net(pre_in,0)
    pre_in=paddle.concat([pre_in,paddle.argmax(outl,-1)[:,-1:]],-1)

print("".join([voc_id[i] for i in pre_in[0]]))


